﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;

namespace DigitClassification
{
    public class NaiveBayesClassifier : BaseClassifier
    {
        public NaiveBayesClassifier(string dataFile, string labelsFile)
        {
            //Build training data
            Training = GetDigitData(dataFile, labelsFile);

            //Count occurances in training data
            CountOccurences();

            //PrintDigits(0,10);
        }

        readonly DigitStats[] digitStatsArr = new DigitStats[10];

        private void CountOccurences()
        {
            //init stats:
            for (int i = 0; i < digitStatsArr.Length; i++)
                digitStatsArr[i] = new DigitStats(Dim);

            //read occurance from training
            foreach (var row in Training)
            {
                DigitStats digitStats = digitStatsArr[row.Label];

                digitStats.Count++; //new digit occurance
                for (int i = 0; i < Dim; i++)
                {
                    for (int j = 0; j < Dim; j++)
                    {
                        if (IsDark(row.Attributes[i, j]))
                            digitStats.CellDarkCount[i, j]++; //new darkcell occurance
                    }
                }
            }
        }

        /// <summary>
        /// classify each image to the label it fit most
        /// return error rate %
        /// </summary>
        /// <param name="sourceName"></param>
        /// <param name="dataFile"></param>
        /// <param name="labelFile"></param>
        /// <param name="kSmoothing"></param>
        /// <returns></returns>
        public double Classify(string sourceName, string dataFile, string labelFile, int kSmoothing)
        {
            var testData = GetDigitData(dataFile, labelFile);
            var errorCount = 0;

            foreach (var item in testData)
            {
                var classifyDigit = Classify(item, kSmoothing);
                var correctDigit = item.Label;

                if (classifyDigit != correctDigit)
                    errorCount++;
            }

            var errorRate = (double)errorCount / testData.Count * 100;
            Console.WriteLine(@"total error on {0} using k={1} is: {2}/{3} ({4}%)",
                sourceName, kSmoothing, errorCount, testData.Count, Math.Round(errorRate, 2));

            return errorRate;
        }

        private int Classify(DigitData item, int kSmoothing)
        {
            var digitProbs = CalcDigitsProbability(item, kSmoothing);
            var mostFitDigit = FindMostFitDigit(digitProbs);

            return mostFitDigit;
        }

        /// <summary>
        /// return array of probabilities for given item
        /// prob[digit] = probability that given item has label=digit
        ///             = Sigma(Log(Pr(f[i,j]=attr[i,j] | digit))) + Log(Pr(digit))
        /// </summary>
        /// <param name="item"></param>
        /// <param name="kSmoothing"></param>
        /// <returns></returns>
        private double[] CalcDigitsProbability(DigitData item, int kSmoothing)
        {
            var probArray = new double[digitStatsArr.Length];
            for (int digit = 0; digit < digitStatsArr.Length; digit++)
            {
                var digitStats = digitStatsArr[digit];
                var prDigit = digitStats.GetProb(Training.Count, kSmoothing); // Pr(digit)

                probArray[digit] = Math.Log(prDigit); //c = Log(Pr(digit))
                for (int i = 0; i < Dim; i++)
                {
                    for (int j = 0; j < Dim; j++)
                    {
                        var prDarkCellGivenDigit = digitStats.GetDarkCellProb(i, j, kSmoothing); // Pr(f[i,j]=1 | C)
                        var prAttributeGivenDigit = IsDark(item.Attributes[i, j])
                            ? prDarkCellGivenDigit
                            : 1 - prDarkCellGivenDigit; // Pr(f[i,j]=attr[i,j] | C)

                        probArray[digit] += Math.Log(prAttributeGivenDigit); //c += Log(Pr(f[i,j]=attr[i,j] | C))
                    }
                }
            }
            return probArray;
        }

        /// <summary>
        /// return digit that has maximum probility
        /// </summary>
        /// <param name="c"></param>
        /// <returns></returns>
        private int FindMostFitDigit(double[] c)
        {
            int digitMostFit = 0;
            var max = c.Max();
            for (int i = 0; i < digitStatsArr.Length; i++)
            {
                if (c[i] < max) continue;

                digitMostFit = i;
                break;
            }
            return digitMostFit;
        }

        /// <summary>
        /// for testing = print digit and label
        /// </summary>
        /// <param name="startIndex"></param>
        /// <param name="endIndex"></param>
        private void PrintDigits(int startIndex, int endIndex)
        {
            for (int i = startIndex; i <= endIndex; i++)
                Console.WriteLine(Training[i].ToString());
        }
    }
}